#include <iostream>
#include <vector>
#include <thread>
#include <chrono>
#include <pthread.h>
#include "readfile.cpp"

// 定义数据结构
struct Centroid {
    double latitude;
    double longitude;
    double averageSignal;
    double range;
};
struct ThreadArgs {
    int start;
    int end;
    std::vector<Point>* data;
    std::vector<Centroid>* centroids;
    std::vector<int>* clusterAssignments;
    std::vector<Centroid>* newCentroids;
    std::vector<int>* counts;

    // 默认构造函数
    ThreadArgs() : start(0), end(0), data(nullptr), centroids(nullptr), clusterAssignments(nullptr), newCentroids(nullptr), counts(nullptr) {}
};

// 计算两点之间的欧几里得距离
double distance(const Point& p, const Centroid& c) {
    return std::sqrt(
        std::pow(p.latitude - c.latitude, 2) +
        std::pow(p.longitude - c.longitude, 2) +
        std::pow(p.averageSignal - c.averageSignal, 2) +
        std::pow(p.range - c.range, 2)
    );
}

// 线程函数：将数据点分配到最近的质心
void* assignPointsToCentroids(void* args) {
    ThreadArgs* threadArgs = static_cast<ThreadArgs*>(args);
    int start = threadArgs->start;
    int end = threadArgs->end;
    std::vector<Point>& data = *(threadArgs->data);
    std::vector<Centroid>& centroids = *(threadArgs->centroids);
    std::vector<int>& clusterAssignments = *(threadArgs->clusterAssignments);

    for (int i = start; i < end; ++i) {
        double minDist = std::numeric_limits<double>::max();
        int bestCluster = 0;
        for (int j = 0; j < centroids.size(); ++j) {
            double dist = distance(data[i], centroids[j]);
            if (dist < minDist) {
                minDist = dist;
                bestCluster = j;
            }
        }
        clusterAssignments[i] = bestCluster;
    }

    return nullptr;
}

// 线程函数：更新质心
void* updateCentroids(void* args) {
    ThreadArgs* threadArgs = static_cast<ThreadArgs*>(args);
    int start = threadArgs->start;
    int end = threadArgs->end;
    std::vector<Point>& data = *(threadArgs->data);
    std::vector<Centroid>& centroids = *(threadArgs->centroids);
    std::vector<int>& clusterAssignments = *(threadArgs->clusterAssignments);
    std::vector<Centroid>& newCentroids = *(threadArgs->newCentroids);
    std::vector<int>& counts = *(threadArgs->counts);

    for (int i = start; i < end; ++i) {
        int cluster = clusterAssignments[i];
        newCentroids[cluster].latitude += data[i].latitude;
        newCentroids[cluster].longitude += data[i].longitude;
        newCentroids[cluster].averageSignal += data[i].averageSignal;
        newCentroids[cluster].range += data[i].range;
        counts[cluster] += 1;
    }

    return nullptr;
}

// 使用多线程实现的K-means算法
void kmeans_pthread(std::vector<Point>& data, int k, int maxIterations, int numThreads) {
    int n = data.size();
    std::vector<Centroid> centroids(k);
    std::srand(std::time(NULL));

    // 随机初始化质心
    for (int i = 0; i < k; ++i) {
        int index = std::rand() % n;
        centroids[i].latitude = data[index].latitude;
        centroids[i].longitude = data[index].longitude;
        centroids[i].averageSignal = data[index].averageSignal;
        centroids[i].range = data[index].range;
    }

    for (int iter = 0; iter < maxIterations; ++iter) {
        // 将数据点分配到最近的质心
        std::vector<int> clusterAssignments(n);
        if (numThreads > 1) {
            std::vector<pthread_t> threads(numThreads);
            std::vector<ThreadArgs> threadArgs(numThreads);

            int chunkSize = n / numThreads;
            for (int i = 0; i < numThreads; ++i) {
                threadArgs[i].start = i * chunkSize;
                threadArgs[i].end = (i == numThreads - 1) ? n : (i + 1) * chunkSize;
                threadArgs[i].data = &data;
                threadArgs[i].centroids = &centroids;
                threadArgs[i].clusterAssignments = &clusterAssignments;
                pthread_create(&threads[i], nullptr, assignPointsToCentroids, &threadArgs[i]);
            }

            for (int i = 0; i < numThreads; ++i) {
                pthread_join(threads[i], nullptr);
            }
        } else {
            for (int i = 0; i < n; ++i) {
                double minDist = std::numeric_limits<double>::max();
                int bestCluster = 0;
                for (int j = 0; j < centroids.size(); ++j) {
                    double dist = distance(data[i], centroids[j]);
                    if (dist < minDist) {
                        minDist = dist;
                        bestCluster = j;
                    }
                }
                clusterAssignments[i] = bestCluster;
            }
        }

        // 临时数组用于并行归约
        std::vector<Centroid> newCentroids(k, {0, 0, 0, 0});
        std::vector<int> counts(k, 0);

        if (numThreads > 1) {
            std::vector<pthread_t> threads(numThreads);
            std::vector<ThreadArgs> threadArgs(numThreads);

            int chunkSize = k / numThreads;
            for (int i = 0; i < numThreads; ++i) {
                threadArgs[i].start = i * chunkSize;
                threadArgs[i].end = (i == numThreads - 1) ? k : (i + 1) * chunkSize;
                threadArgs[i].data = &data;
                threadArgs[i].centroids = &centroids;
                threadArgs[i].clusterAssignments = &clusterAssignments;
                threadArgs[i].newCentroids = &newCentroids;
                threadArgs[i].counts = &counts;
                pthread_create(&threads[i], nullptr, updateCentroids, &threadArgs[i]);
            }

            for (int i = 0; i < numThreads; ++i) {
                pthread_join(threads[i], nullptr);
            }
        } else {
            for (int i = 0; i < n; ++i) {
                int cluster = clusterAssignments[i];
                newCentroids[cluster].latitude += data[i].latitude;
                newCentroids[cluster].longitude += data[i].longitude;
                newCentroids[cluster].averageSignal += data[i].averageSignal;
                newCentroids[cluster].range += data[i].range;
                counts[cluster] += 1;
            }
        }

        // 更新质心
        for (int j = 0; j < k; ++j) {
            if (counts[j] > 0) {
                centroids[j].latitude = newCentroids[j].latitude / counts[j];
                centroids[j].longitude = newCentroids[j].longitude / counts[j];
                centroids[j].averageSignal = newCentroids[j].averageSignal / counts[j];
                centroids[j].range = newCentroids[j].range / counts[j];
            }
        }
    }
}
